#!/usr/bin/env python3
"""
批量可视化轨迹脚本

功能：
1. 读取完整轨迹数据集（vae_full_trajectories_*.npz）
2. 每5条轨迹画在一张图上（同时显示ego和opponent）
3. 保存到单独的文件目录

数据格式：
- trajectories: [n_episodes, n_steps, 12] - 完整轨迹
- ego_goals: [n_episodes, 3]
- opponent_goals: [n_episodes, 3]
"""

import os
import sys
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from glob import glob
from typing import Optional

# 添加路径
repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
sys.path.insert(0, repo_root)


def plot_batch_trajectories_3d(
    trajectories: np.ndarray,
    ego_goals: np.ndarray,
    opponent_goals: np.ndarray,
    episode_indices: list,
    batch_id: int,
    save_path: Optional[str] = None
):
    """
    在一张图上绘制多条轨迹（每5条）
    
    Args:
        trajectories: [n_episodes, n_steps, 12] 完整轨迹
        ego_goals: [n_episodes, 3] Ego目标位置
        opponent_goals: [n_episodes, 3] Opponent目标位置
        episode_indices: 要绘制的episode索引列表
        batch_id: 批次编号
        save_path: 保存路径
    """
    # 创建3D图
    fig = plt.figure(figsize=(20, 6))
    
    # 定义颜色列表（为每条轨迹使用不同颜色）
    colors_ego = ['blue', 'cyan', 'green', 'orange', 'purple']
    colors_opp = ['red', 'magenta', 'brown', 'pink', 'salmon']
    
    # 子图1: 3D轨迹
    ax1 = fig.add_subplot(131, projection='3d')
    
    # 子图2: XY平面
    ax2 = fig.add_subplot(132)
    
    # 子图3: 高度随时间变化
    ax3 = fig.add_subplot(133)
    
    # 绘制每条轨迹
    for idx, episode_id in enumerate(episode_indices):
        if episode_id >= len(trajectories):
            continue
            
        trajectory = trajectories[episode_id]  # [n_steps, 12]
        ego_goal = ego_goals[episode_id]  # [3]
        opponent_goal = opponent_goals[episode_id]  # [3]
        
        # 提取轨迹
        ego_traj = trajectory[:, 0:3]      # [n_steps, 3] - px, py, pz
        opponent_traj = trajectory[:, 6:9]  # [n_steps, 3] - px, py, pz
        
        color_ego = colors_ego[idx % len(colors_ego)]
        color_opp = colors_opp[idx % len(colors_opp)]
        alpha = 0.7
        
        # 3D轨迹
        ax1.plot(ego_traj[:, 0], ego_traj[:, 1], ego_traj[:, 2], 
                 color=color_ego, linewidth=1.5, label=f'Ego-{episode_id}', alpha=alpha)
        ax1.plot(opponent_traj[:, 0], opponent_traj[:, 1], opponent_traj[:, 2], 
                 color=color_opp, linewidth=1.5, label=f'Opp-{episode_id}', alpha=alpha, linestyle='--')
        
        # 起点（只在第一条轨迹时显示）
        if idx == 0:
            ax1.scatter(ego_traj[0, 0], ego_traj[0, 1], ego_traj[0, 2], 
                       c=color_ego, marker='o', s=80, alpha=0.8)
            ax1.scatter(opponent_traj[0, 0], opponent_traj[0, 1], opponent_traj[0, 2], 
                       c=color_opp, marker='o', s=80, alpha=0.8)
        
        # 目标点
        ax1.scatter(ego_goal[0], ego_goal[1], ego_goal[2], 
                   c=color_ego, marker='*', s=200, alpha=0.8, edgecolors='black', linewidths=0.5)
        ax1.scatter(opponent_goal[0], opponent_goal[1], opponent_goal[2], 
                   c=color_opp, marker='*', s=200, alpha=0.8, edgecolors='black', linewidths=0.5)
        
        # XY平面
        ax2.plot(ego_traj[:, 0], ego_traj[:, 1], 
                color=color_ego, linewidth=1.5, label=f'Ego-{episode_id}', alpha=alpha)
        ax2.plot(opponent_traj[:, 0], opponent_traj[:, 1], 
                color=color_opp, linewidth=1.5, label=f'Opp-{episode_id}', alpha=alpha, linestyle='--')
        
        if idx == 0:
            ax2.scatter(ego_traj[0, 0], ego_traj[0, 1], c=color_ego, marker='o', s=80, alpha=0.8)
            ax2.scatter(opponent_traj[0, 0], opponent_traj[0, 1], c=color_opp, marker='o', s=80, alpha=0.8)
        
        ax2.scatter(ego_goal[0], ego_goal[1], c=color_ego, marker='*', s=200, 
                   alpha=0.8, edgecolors='black', linewidths=0.5)
        ax2.scatter(opponent_goal[0], opponent_goal[1], c=color_opp, marker='*', s=200, 
                   alpha=0.8, edgecolors='black', linewidths=0.5)
        
        # 高度随时间变化
        time_steps = np.arange(len(trajectory)) * 0.1  # dt=0.1
        ax3.plot(time_steps, ego_traj[:, 2], 
                color=color_ego, linewidth=1.5, label=f'Ego-{episode_id}', alpha=alpha)
        ax3.plot(time_steps, opponent_traj[:, 2], 
                color=color_opp, linewidth=1.5, label=f'Opp-{episode_id}', alpha=alpha, linestyle='--')
        
        if idx == 0:
            ax3.axhline(y=ego_goal[2], color=color_ego, linestyle=':', alpha=0.5, linewidth=1)
            ax3.axhline(y=opponent_goal[2], color=color_opp, linestyle=':', alpha=0.5, linewidth=1)
    
    # 设置标签和标题
    ax1.set_xlabel('X (m)', fontsize=10)
    ax1.set_ylabel('Y (m)', fontsize=10)
    ax1.set_zlabel('Z (m)', fontsize=10)
    ax1.set_title(f'Batch {batch_id}: 3D Trajectories (Episodes {episode_indices[0]}-{episode_indices[-1]})', fontsize=12)
    ax1.legend(loc='upper left', fontsize=8, ncol=2)
    ax1.grid(True, alpha=0.3)
    
    ax2.set_xlabel('X (m)', fontsize=10)
    ax2.set_ylabel('Y (m)', fontsize=10)
    ax2.set_title('XY Plane View', fontsize=12)
    ax2.legend(loc='upper left', fontsize=8, ncol=2)
    ax2.grid(True, alpha=0.3)
    ax2.axis('equal')
    
    ax3.set_xlabel('Time (s)', fontsize=10)
    ax3.set_ylabel('Altitude (m)', fontsize=10)
    ax3.set_title('Altitude Profile', fontsize=12)
    ax3.legend(loc='upper left', fontsize=8, ncol=2)
    ax3.grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"  ✓ 批次 {batch_id} 已保存: {save_path}")
    
    plt.close()  # 关闭图形以释放内存


def visualize_all_trajectories(
    dataset_path: str,
    output_dir: Optional[str] = None,
    trajectories_per_batch: int = 5
):
    """
    批量可视化所有轨迹
    
    Args:
        dataset_path: 完整轨迹数据集路径
        output_dir: 输出目录（默认为数据集同目录下的visualizations文件夹）
        trajectories_per_batch: 每批次的轨迹数量
    """
    print(f"\n{'='*70}")
    print(f"📊 批量轨迹可视化")
    print(f"{'='*70}\n")
    
    # 加载数据
    print(f"[1] 加载数据集: {dataset_path}")
    data = np.load(dataset_path)
    trajectories = data["trajectories"]  # [n_episodes, n_steps, 12]
    ego_goals = data["ego_goals"]  # [n_episodes, 3]
    opponent_goals = data["opponent_goals"]  # [n_episodes, 3]
    
    n_episodes = len(trajectories)
    print(f"  ✓ 加载完成: {n_episodes} 条轨迹")
    print(f"  ✓ 轨迹形状: {trajectories.shape}")
    print(f"  ✓ Ego目标形状: {ego_goals.shape}")
    print(f"  ✓ Opponent目标形状: {opponent_goals.shape}")
    
    # 创建输出目录
    if output_dir is None:
        dataset_dir = os.path.dirname(dataset_path)
        output_dir = os.path.join(dataset_dir, "trajectory_visualizations")
    
    os.makedirs(output_dir, exist_ok=True)
    print(f"\n[2] 输出目录: {output_dir}")
    
    # 计算批次数量
    n_batches = (n_episodes + trajectories_per_batch - 1) // trajectories_per_batch
    print(f"\n[3] 开始批量可视化...")
    print(f"  ✓ 每批次轨迹数: {trajectories_per_batch}")
    print(f"  ✓ 总批次数: {n_batches}")
    print(f"\n{'='*70}\n")
    
    # 批量绘制
    for batch_id in range(n_batches):
        start_idx = batch_id * trajectories_per_batch
        end_idx = min(start_idx + trajectories_per_batch, n_episodes)
        episode_indices = list(range(start_idx, end_idx))
        
        print(f"[批次 {batch_id+1}/{n_batches}] Episodes {start_idx}-{end_idx-1} ({len(episode_indices)}条)")
        
        # 生成保存路径
        save_path = os.path.join(output_dir, f"batch_{batch_id+1:02d}_episodes_{start_idx:03d}_{end_idx-1:03d}.png")
        
        # 绘制
        plot_batch_trajectories_3d(
            trajectories=trajectories,
            ego_goals=ego_goals,
            opponent_goals=opponent_goals,
            episode_indices=episode_indices,
            batch_id=batch_id+1,
            save_path=save_path
        )
    
    print(f"\n{'='*70}")
    print(f"✅ 批量可视化完成!")
    print(f"{'='*70}")
    print(f"\n输出统计:")
    print(f"  总轨迹数: {n_episodes}")
    print(f"  总批次数: {n_batches}")
    print(f"  输出目录: {output_dir}")
    print(f"  文件格式: batch_XX_episodes_XXX_XXX.png")
    print(f"\n")


def find_latest_dataset(dataset_dir: str) -> Optional[str]:
    """
    查找最新的完整轨迹数据集文件
    
    Args:
        dataset_dir: 数据集目录
        
    Returns:
        最新数据集文件路径，如果不存在则返回None
    """
    pattern = os.path.join(dataset_dir, "vae_full_trajectories_*.npz")
    files = glob(pattern)
    
    if not files:
        return None
    
    # 按修改时间排序，返回最新的
    files.sort(key=os.path.getmtime, reverse=True)
    return files[0]


def main():
    """
    主函数
    """
    import argparse
    
    parser = argparse.ArgumentParser(description="批量可视化轨迹数据集")
    parser.add_argument(
        "--dataset",
        type=str,
        default=None,
        help="完整轨迹数据集路径（.npz文件）。如果不指定，将自动查找最新的数据集。"
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="输出目录。如果不指定，将在数据集同目录下创建trajectory_visualizations文件夹。"
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=5,
        help="每批次的轨迹数量（默认：5）"
    )
    
    args = parser.parse_args()
    
    # 确定数据集路径
    if args.dataset:
        dataset_path = args.dataset
    else:
        # 自动查找最新的数据集
        script_dir = os.path.dirname(os.path.abspath(__file__))
        dataset_dir = os.path.join(script_dir, "vaedataset")
        dataset_path = find_latest_dataset(dataset_dir)
        
        if dataset_path is None:
            print(f"❌ 错误: 在 {dataset_dir} 中未找到完整轨迹数据集文件")
            print(f"   请使用 --dataset 参数指定数据集路径")
            return
        else:
            print(f"📁 自动找到最新数据集: {dataset_path}")
    
    if not os.path.exists(dataset_path):
        print(f"❌ 错误: 数据集文件不存在: {dataset_path}")
        return
    
    # 执行可视化
    visualize_all_trajectories(
        dataset_path=dataset_path,
        output_dir=args.output,
        trajectories_per_batch=args.batch_size
    )


if __name__ == "__main__":
    main()

